package com.zdp.agent;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.Objects;
import java.util.Optional;

/**
 * 为类添加方法
 * @author Zhengdp
 * @date 2023年02月15日 10:14
 */
public class AddMethodClassFileTransformer implements ClassFileTransformer {
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        // 将自定义的代码的所有方法都插入输出语句
        System.out.println("正在加载类："+ className);
        if(className == null || !className.startsWith("com/zdp")){
            return classfileBuffer;
        }
        CtClass cl = null;
        try {
            ClassPool classPool = ClassPool.getDefault();
            cl = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
            CtMethod[] declaredMethods = cl.getDeclaredMethods();
            if(Objects.nonNull(declaredMethods) && declaredMethods.length != 0) {
                for(CtMethod ctMethod : declaredMethods) {
                    System.out.println("获取方法名称："+ ctMethod.getName());
                    if(!"main".equals(ctMethod.getName())){
                        ctMethod.insertBefore("System.out.println(\" 动态插入的打印语句 \");");
                        ctMethod.insertAfter("System.out.println($_);");
                    }
                }
            }
            byte[] transformed = cl.toBytecode();
            return transformed;
        }catch (Exception e){
            e.printStackTrace();
        }
        return classfileBuffer;
    }
}
